# compare with ~/xchg/Match/cem/a1lalonde_nocaliper.R, except with estimand=ATT

#add in machine learning algorithms, from ~/xchg/genmatch_paper/LeeLesslerStaurt/R/simsA1.nobs1000.R

#the truth for lalonde, regression model

#compare with lapo:~/xchg/Match/GenMatch/mc1/lalonde3matching1.Rout

source("../LeeLesslerStaurt/R/pscore_estimation_core1.R")
library(Matching)
data(lalonde)

sims <- 1000
estimand <- "ATT"
obs <- nrow(lalonde)

#Run Seed
set.seed(47191)

attach(lalonde)
propensity  <- glm(treat~ I(age^2) + I(educ^2) + black +
                   hisp + married + nodegr + I(re74^2) + I(re75^2) +
                   u74 + u75, family=binomial, data=lalonde)

X <- cbind(rep(1, length(treat)),
           propensity$linear.pred,
           I(log(age)^2),
           I(log(educ)^2),
           I(log(re74+0.01)^2),
           I(log(re75+0.01)^2))

propensity.coeffs <- propensity$coeff
propensity.coeffs <- as.matrix(c(
                                 1+00,  #(Intercept)        
                                 .5,    #linear.pred
                                 .01, #age
                                 -.3,  #educ
                                 -0.01, #I(re74^2)          
                                 0.01   #I(re75^2)          
                                 ))

mu = X %*% propensity.coeffs
Tr.pred <- exp(mu)/(1+exp(mu))
#print(summary(Tr.pred))

TreatmentEffect <- 1000
TreatmentReal <- matrix(nrow=obs, ncol=1)

#return objects
raw <- matrix(nrow=sims)
correct.est.bias <- matrix(nrow=sims)
correct.est.nobias <- matrix(nrow=sims)
est.bias <- matrix(nrow=sims)
est.nobias <- matrix(nrow=sims)
iv.est.bias <- matrix(nrow=sims)
iv.est.nobias <- matrix(nrow=sims)
ps.est.bias <- matrix(nrow=sims)
ps.est.nobias <- matrix(nrow=sims)
psmh.est.bias <- matrix(nrow=sims)
psmh.est.nobias <- matrix(nrow=sims)


boosting.pscore <- vector("list",sims)
boosting.desc <- vector("list",sims)
boosting.est <- matrix(0, nrow=sims, ncol=1)
boosting.nSE <- matrix(0, nrow=sims, ncol=1)
boosting.aiSE <- matrix(0, nrow=sims, ncol=1)

forest.pscore <- vector("list",sims)
forest.est <- matrix(0, nrow=sims, ncol=1)
forest.nSE <- matrix(0, nrow=sims, ncol=1)
forest.aiSE <- matrix(0, nrow=sims, ncol=1)

nnet.pscore <- vector("list",sims)
nnet.est <- matrix(0, nrow=sims, ncol=1)
nnet.nSE <- matrix(0, nrow=sims, ncol=1)
nnet.aiSE <- matrix(0, nrow=sims, ncol=1)


results <- function(s)
  {

    
    cat("\n")
    
  }

#load simulation data
load(file="RData.mc2.data1")

count <- 0
for (s in 1:sims)
  {

    cat("\nS:", s, "\n")
    count <-  count+1

    TreatmentReal <-  sims.data[[s]][,2]
    Y <- sims.data[[s]][,1]

    raw[s] = mean(Y[TreatmentReal==1])-mean(Y[TreatmentReal==0]) - TreatmentEffect
    cat("raw[",s,"]:", raw[s], "\n")
    cat("RAW bias",s,":",mean(raw[1:s]), "\n")
    cat("RAW RMSE",s,":", sqrt(mean(raw[1:s]^2)), "\n")        

    Xtmp = cbind(age, educ, black, hisp, married, nodegr, re74, re75, u74, u75)
    n1 <- nnet(x=Xtmp, y=TreatmentReal, size=40, trace=FALSE)
    nnet.pscore[[s]] <- n1$fitted.values
    
    m1 <- Match(Y=Y, Tr=TreatmentReal, X=Xtmp, estimand="ATT", M=1)
    nnet.est[s] <-  m1$est-TreatmentEffect
    nnet.nSE[s] <- m1$se.standard
    nnet.aiSE[s] <- m1$se

    cat("\nnnet",s,":",nnet.est[s],"\n")
    if(s>1)
      {
        cat("nnet mean",s,":",mean(nnet.est[1:s]),"\n")
        rmse <- sqrt(mean((nnet.est[1:s])^2))
        cat("nnet RMSE",s,":",rmse,"\n")
      }

    #Random Forest
    ps1 <- 1
    while (max(ps1)> .999) {
#      gc(verbose=TRUE)
      fps.model <- randomForest(as.factor(TreatmentReal) ~ age+educ+black+hisp+married+nodegr+re74+re75+u74+u75)
      ps1 <- fps.model$votes[, 2]
    }
    
    forest.pscore[[s]] <- ps1

    m1 <- Match(Y=Y, Tr=TreatmentReal, X=ps1, estimand="ATT", M=1)
    forest.est[s] <-  m1$est-TreatmentEffect
    forest.nSE[s] <- m1$se.standard
    forest.aiSE[s] <- m1$se

    cat("\nforest",s,":",forest.est[s],"\n")
    if(s>1)
      {
        cat("forest mean",s,":",mean(forest.est[1:s]),"\n")
        rmse <- sqrt(mean((forest.est[1:s])^2))
        cat("forest RMSE",s,":",rmse,"\n")
      }

    #boosting!
    sdta <- as.data.frame(cbind(TreatmentReal,age,educ,black,hisp,married,nodegr,re74,re75,u74,u75))
    ps.model <- ps(TreatmentReal ~ age+educ+black+hisp+married+nodegr+re74+re75+u74+u75,
                   data=sdta, title="none", stop.method=stop.methods["ks.stat.mean"],
                   plots="none", pdf.plots=F, n.trees=20000, interaction.depth=4,
                   shrinkage=0.0005, verbose=FALSE)

    boosting.pscore[[s]] <- ps.model$ps
    boosting.desc[[s]] <- ps.model$desc

    m1 <- Match(Y=Y, Tr=TreatmentReal, X=ps.model$ps, estimand="ATT", M=1)
    boosting.est[s] <-  m1$est-TreatmentEffect
    boosting.nSE[s] <- m1$se.standard
    boosting.aiSE[s] <- m1$se

    cat("\nboosting",s,":",boosting.est[s],"\n")
    if(s>1)
      {
        cat("boosting mean",s,":",mean(boosting.est[1:s]),"\n")
        rmse <- sqrt(mean((boosting.est[1:s])^2))
        cat("boosting RMSE",s,":",rmse,"\n")
      }    
    
    if (count==100)
      {
        results(s)
        count <-  0
        save.image(file="RDATA.mc2.machine1")
      }
  } #end of sims

results(s)

save.image(file="RDATA.mc2.machine1")
